import pandas as pd
import os,glob,math
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.lines as mlines
import matplotlib.dates as mdates
import cartopy.crs as ccrs
from cartopy.mpl.ticker import LongitudeFormatter,LatitudeFormatter
import cartopy.feature as cfeature
import shapely.geometry as sgeom
from datetime import datetime,timedelta
from tqdm import tqdm
from global_land_mask import globe
from gsj_typhoon import tydat,see,count_rapidgrow,tydat_CMA,average_datetime,split_str_id,load_land_polygons,detect_landfall
from geopy.distance import geodesic
import matplotlib.ticker as ticker
import seaborn as sns
from typlot.config.global_config import *


names = ['mojie_28','dusurui_16','gaemi_09','haikui_38','kangni_54','shantuo_44','saola_25','koinu_49']
# names = ['koinu_49']
tynames,tyids = split_str_id(names)
draw_obs_opt = True
draw_opt = False

all_ini_time_mode = ['00_12','12','00']
all_obs_baseline = ['RI','land']   # RI  land




def add_red_dots(obs_arr: np.ndarray , ini_time_mode : str ) :
    ''' 补全观测红点，策略为先y后x循环
    obs_arr: 拓展到正确形状之后的数组。（指最上条红线恰好过原点）
    '''
    assert ini_time_mode in ['00','12','00_12'],'Wrong ini_time_mode value'
    assert obs_arr.ndim==2, 'Wrong obs_arr shape, ndim should be 2'
    ny,nx = obs_arr.shape
    obs_arr = obs_arr[::-1]  # 代码逻辑更直观一些
    dx = 2 if ini_time_mode=='00_12' else 1
    for j in range(ny-1,-1,-1):
        for i in range(nx):
            try:
                if obs_arr[j+1,i+dx] ==1 :
                    obs_arr[j,i]=1
            except: 
                pass #懒得处理数组越界问题了
    return obs_arr[::-1]



for ini_time_mode in all_ini_time_mode:
    for obs_baseline in all_obs_baseline:  
        print('Setting:',ini_time_mode,obs_baseline)
        draw_data_list = []
        obs_list = []
        
        for ty,tyid in zip(tynames,tyids):
            ''' init '''
            RIstd = 7
            dt = timedelta(days=1)
            directory = os.path.join(global_ensdir,f'{ty}_{tyid}')
            dates_name = sorted(os.listdir(directory))
            dates_name = [i for i in dates_name if i[-2:] in ini_time_mode]
            dates_paths = [os.path.join(directory,date) for date in dates_name ]
            obs_path = os.path.join(global_obsdir,f'{ty}_CMAobs.txt')
            pic_savepath = os.path.join(global_picdir,f'{obs_baseline}_heatmap',ini_time_mode,'81')
            os.makedirs(pic_savepath,exist_ok=True)
            
            ''' calc 2d numbers '''
            # 遍历所有起报时间
            max_delta_time = timedelta(days=0)
            for j,i in enumerate(dates_paths):
                mem_paths = glob.glob( os.path.join(i,"TRACK_ID_*") )
                # sorted_mem_paths = sorted(mem_paths, key=lambda x: int(x.split('_')[-1]))
                start_time = datetime.strptime(dates_name[j],"%Y%m%d%H")
                # 遍历所有成员
                for single_mem_path in mem_paths:
                    single_mem = tydat(single_mem_path,RIstd) 
                    delta_time = single_mem.time[-1] - start_time
                    if delta_time>max_delta_time:
                        max_delta_time = delta_time
            max_dt_num = math.ceil(max_delta_time/dt) 
            
            ''' 求各起报时间的图 '''
            tyobs = tydat_CMA(obs_path)
            rcd_obs = count_rapidgrow(RIstd, tyobs.umax, tyobs.time)
            num_result = np.zeros(( max_dt_num,len(dates_name) ))
            dis_result = np.zeros(( max_dt_num,len(dates_name) ))
            dis_var_result = np.zeros(( max_dt_num,len(dates_name) ))
            
            
            for i in tqdm( range(len(dates_paths)), desc=f'processing  {ty}'):
                '''读取这个起报时间的所有集合'''
                mem_paths = glob.glob( os.path.join(dates_paths[i],"TRACK_ID_*") )
                sorted_mem_paths = sorted(mem_paths, key=lambda x: int(x.split('_')[-1]))
                ty_mem = [tydat(i,RIstd) for i in sorted_mem_paths]
                start_dt = datetime.strptime( dates_name[i],"%Y%m%d%H" )
                for j in range(max_dt_num):
                    t0 = j
                    t1 = j+1
                    t0_dt = start_dt + timedelta(days=t0)
                    t1_dt = start_dt + timedelta(days=t1)
                    '''如果观测在这两天没有RI，就直接跳过'''
                    idx_obs_2d = (tyobs.time>=t0_dt)&(tyobs.time<t1_dt)
                    idx_obs_2d_RI = (rcd_obs[idx_obs_2d]==1)
                    if np.sum(rcd_obs[idx_obs_2d]) ==0.:
                        ''' only record numbers here '''
                        for k in range(len(ty_mem)):
                            '''cycle through all members : 如果这个成员在这两天没有RI,就跳过'''
                            idx_mem_2d    = (ty_mem[k].time>=t0_dt)&(ty_mem[k].time<t1_dt)
                            idx_mem_2d_RI = (ty_mem[k].num_rapidgrow()[idx_mem_2d]==1)
                            num_result[j,i] += np.sum(idx_mem_2d_RI)
                        continue
                    else:
                        '''如果2d内有多个RI,取平均'''
                        time_obs_ave = average_datetime(tyobs.time[idx_obs_2d][idx_obs_2d_RI])
                        lon_obs_ave  = np.mean(tyobs.lon[idx_obs_2d][idx_obs_2d_RI])
                        lat_obs_ave  = np.mean(tyobs.lat[idx_obs_2d][idx_obs_2d_RI])
                        '''开始计算'''
                        num = 0 
                        dis_list = [] 
        
                        for k in range(len(ty_mem)):
                            '''cycle through all members : 如果这个成员在这两天没有RI,就跳过'''
                            idx_mem_2d    = (ty_mem[k].time>=t0_dt)&(ty_mem[k].time<t1_dt)
                            idx_mem_2d_RI = (ty_mem[k].num_rapidgrow()[idx_mem_2d]==1)
                            time_mem = ty_mem[k].time[idx_mem_2d][idx_mem_2d_RI]
                            lat_mem  = ty_mem[k].lat[idx_mem_2d][idx_mem_2d_RI]
                            lon_mem  = ty_mem[k].lon[idx_mem_2d][idx_mem_2d_RI]
                            if np.sum(idx_mem_2d_RI) ==0.:
                                continue
                            else:
                                num+=np.sum(idx_mem_2d_RI)
                                for l in range(len(lon_mem)):
                                    dis_list.append(geodesic((lat_obs_ave,lon_obs_ave),(lat_mem[l],lon_mem[l])).kilometers)
        
                    if num==0:
                        continue
                    num_result[j,i] = num
                    dis_result[j,i] = np.array(dis_list).mean()
                    dis_var_result[j,i] = np.array(dis_list).std()
            draw_data_list.append(dis_result[:15])
        
                    
            if draw_obs_opt == True : 
                obs = tydat_CMA(obs_path) 
                draw_data_obs = np.zeros( (len(dates_name),max_dt_num),dtype=int )
                tt,dd = draw_data_obs.shape
                if obs_baseline == 'RI':
                   record_obs = count_rapidgrow(RIstd,obs.umax, obs.time)
                   for t in range(tt):
                       #预报时效
                       for d in range(dd):
                           try:
                               start_time = datetime.strptime(dates_name[t],"%Y%m%d%H")
                               cal_time = start_time + d*timedelta(days=1)
                               calend_time = cal_time + timedelta(days=1)
                               # 对齐
                               trange = (obs.time>=cal_time) & (obs.time<calend_time)
                               draw_data_obs[t][d] = 1 if np.sum(record_obs[trange])>0 else 0 
                           except Exception as e:
                               print(e)       
                if obs_baseline == 'land':
                    # get  land  time
                    land_polys = load_land_polygons(os.path.join(global_picdir,'China','bou1_4p.shp'),os.path.join(global_picdir,'China','bou1_4p.dbf'))
                    flags,id1 = detect_landfall(obs.lat, obs.lon, land_polys)
                    if id1 is None:
                        pass
                    else:
                        land_time = obs.time[id1]
                        # identify  position
                        for t in range(tt):
                            for d in range(dd):
                                start_time = datetime.strptime(dates_name[t],"%Y%m%d%H")
                                cal_time = start_time + d*timedelta(days=1)
                                calend_time = cal_time + timedelta(days=1)
                                trange = (obs.time>=cal_time) & (obs.time<calend_time)
                                # 分类trange； 即有无登陆点
                                draw_data_obs[t][d] = 1 if land_time in obs.time[trange] else 0
            obs_list.append(draw_data_obs.T[:15])
        
        
        ### 开始画图
        mode = 'upall'
                    
        
        if mode=='upall':
            # obs_list  draw_data_list
            assert len(draw_data_list)==len(obs_list), 'obs and ensemble not corespond'
            new_obs_list,new_mem_list = [],[]
            for n in range(len(draw_data_list)):
                print('processing number',n)
                obs = obs_list[n]
                mem = draw_data_list[n]
                # 无红点情况
                if sum(sum(obs)) == 0:
                    new_mem_list.append(mem)
                    new_obs_list.append(obs)
                # 红点在左边界
                if sum(obs[:,0])>0:
                    topindex = np.min(np.where(obs[:,0]==1))
                    # 分两种红点类型
                    if obs[topindex,1]==0: # 00/12 总是这个情况
                        day_ahead = obs.shape[0]-1-topindex
                        x_compensate = 1+day_ahead*2 if ini_time_mode=='00_12' else day_ahead
                        padded_obs = np.pad(obs, pad_width=((0, 0), (x_compensate, 0)),mode='constant', constant_values=0)
                        padded_mem = np.pad(mem, pad_width=((0, 0), (x_compensate, 0)),mode='constant', constant_values=0)
                        new_obs_list.append(add_red_dots(padded_obs,ini_time_mode))
                        new_mem_list.append(padded_mem)
                    elif obs[topindex,1]==1:  
                        day_ahead = obs.shape[0]-1-topindex
                        x_compensate = day_ahead*2 if ini_time_mode=='00_12' else day_ahead
                        padded_obs = np.pad(obs, pad_width=((0, 0), (x_compensate, 0)),mode='constant', constant_values=0)
                        padded_mem = np.pad(mem, pad_width=((0, 0), (x_compensate, 0)),mode='constant', constant_values=0)
                        new_obs_list.append(add_red_dots(padded_obs,ini_time_mode))
                        new_mem_list.append(padded_mem)
                # 红点下边界
                elif sum(obs[:,0])==0 and sum(sum(obs)!=0):
                    leftindex = np.min(np.where(obs[-1,:]==1))
                    new_mem_list.append(mem[:,leftindex:])
                    new_obs_list.append(obs[:,leftindex:])
            
            # 合并8个补全后台风数组
            max_xlength = max([mem.shape[1] for mem in new_mem_list])
            new_mem_list = [np.pad(mem, pad_width=((0, 0), (0, max_xlength-mem.shape[1])),mode='constant', constant_values=0) 
                            for mem in new_mem_list]
            new_obs_list = [np.pad(obs, pad_width=((0, 0), (0, max_xlength-obs.shape[1])),mode='constant', constant_values=0) 
                            for obs in new_obs_list]
            draw_data = np.array(new_mem_list).sum(axis=0)
            
            
            
                        
                        
            # draw heatmap of RI number 
            draw_data = np.array(new_mem_list).sum(axis=0)
            
            # 🎯 步骤1：先找到最大的maxindex
            maxindex = 0
            if draw_obs_opt == True and len(new_obs_list) > 0:
                for draw_data_obs in new_obs_list:
                    points = np.where(draw_data_obs == 1)
                    if len(points[1]) > 0:  # 确保有观测点
                        current_max = points[1].max()
                        maxindex = max(maxindex, current_max)
            
            print(f"找到的maxindex: {maxindex}")
            
            # 🎯 步骤2：根据maxindex裁剪数据
            if maxindex > 0:
                # 裁剪到maxindex+1列（包含maxindex）
                draw_data_cropped = draw_data[:, :maxindex+1]
                # 同时裁剪观测数据
                new_obs_list_cropped = [obs[:, :maxindex+1] for obs in new_obs_list]
            else:
                # 如果没有观测点，寻找数据的自然边界
                non_zero_cols = np.any(draw_data > 0, axis=0)
                if np.any(non_zero_cols):
                    last_data_col = np.max(np.where(non_zero_cols)[0])
                    crop_end = min(last_data_col + 3, draw_data.shape[1])
                else:
                    crop_end = min(20, draw_data.shape[1])  # 默认20列
                
                draw_data_cropped = draw_data[:, :crop_end]
                new_obs_list_cropped = [obs[:, :crop_end] for obs in new_obs_list]
            
            print(f"裁剪后数据形状: {draw_data_cropped.shape}")
            
            
            if False:
                # 🎯 步骤3：创建图形并绘制热图
                plt.figure(figsize=(15, 6), dpi=1000)
                ax = sns.heatmap(draw_data_cropped/len(names), 
                                 annot=True,
                                 annot_kws={'size': 4}, 
                                 fmt=".1f", 
                                 cmap="YlGnBu", 
                                 cbar_kws={'label': 'sum average RI location gap between obs and ensemble members Units: KM'}, 
                                 linewidths=0.5,
                                 vmin=0, 
                                 vmax=500)
                
                # 🎯 步骤4：设置标签和格式
                ax.set_ylabel("Lead Time", fontsize=12)
                ax.set_xlabel("Initialization Time Minus First Initialization Time: h", fontsize=12)
                
                # 基于裁剪后的数据设置x轴标签
                dt = 12 if ini_time_mode=='00_12' else 24
                x_labels = [str(i * dt) for i in range(draw_data_cropped.shape[1])]
                ax.set_xticklabels(x_labels)
                
                plt.yticks(rotation=0)
                plt.xticks(rotation=90)
                
                # 确保y轴标签数量匹配
                ylabels = [str(i)+"d" for i in range(1, min(16, draw_data_cropped.shape[0] + 1))]
                ax.set_yticklabels(ylabels[:draw_data_cropped.shape[0]])
                
                ax.set_title(f'average distance between RIobs and RImems of average 8 typhoons forecast_ini_mode:{ini_time_mode}')
                
                # 设置colorbar
                cbar = ax.collections[0].colorbar
                cbar.locator = ticker.MaxNLocator(integer=True)
                cbar.update_ticks()
                
                # 🎯 步骤5：在热图上绘制观测点
                if draw_obs_opt == True:
                    for draw_data_obs in new_obs_list_cropped:
                        points = np.where(draw_data_obs == 1)
                        if len(points[1]) > 0:
                            y_coords = points[0] + 0.5
                            x_coords = points[1] + 0.15
                            ax.scatter(x_coords, y_coords, color='red', marker='o', s=10, zorder=5)
                
                # 保存和关闭
                if draw_opt==False:
                    plt.tight_layout()
                    plt.savefig(f"{pic_savepath}\\start_day_distance_81.png", dpi=900, bbox_inches='tight')
                    plt.close()
                else:
                    plt.show()
            
            else :
                '''绘制 x-lead_time  y-init_time 图；最棘手的是红点逻辑全变 '''
                plt.figure(figsize=(15, 6))
                arr = draw_data_cropped.T
                obsarr = np.array(new_obs_list_cropped)[0].T
                ax = sns.heatmap(arr /len(names), 
                                 annot=False,
                                 annot_kws={'size': 6}, 
                                 # 控制方格内数字显示的格式
                                 fmt=".0f", 
                                 cmap="YlGnBu", 
                                 # format设置成整数，因为数组中的数值虽然是整数但类型是float，cbar刻度带小数点
                                 cbar_kws={'label': 'Unit: km','format':'%d'}, 
                                 linewidths=0.5,
                                 # 设置cbar刻度范围
                                 vmin=0, 
                                 vmax=500)
                # description
                ax.set_title(f'Mean Distance between Observed and Forecast RI [{ini_time_mode} {obs_baseline}]')
                ax.set_xlabel('Lead Time')
                ax.set_ylabel('Relative to Unified RI Start Time (8 TCs)')
                
                # tick
                ax.set_xticklabels([f'{i+1}d' for i in range(arr.shape[1])])
                if ini_time_mode=='00_12':
                    yticks      = [i+0.5 for i in range(arr.shape[0]) if (i*12)%48==0]
                    yticklabels = [f"{(i*12)//24-14}d" for i in range(arr.shape[0]) if (i*12)%48==0]
                    ax.set_yticks(yticks)
                    ax.set_yticklabels(yticklabels, rotation = 0)
                else:
                    yticks      = [i+0.5 for i in range(arr.shape[0]) if (i*24)%48==0]
                    yticklabels = [f"{(i*24)//24-14}d" for i in range(arr.shape[0]) if (i*24)%48==0]
                    ax.set_yticks(yticks)
                    ax.set_yticklabels(yticklabels, rotation = 0)
                    
                if draw_obs_opt == True:
                    for draw_data_obs in new_obs_list_cropped:
                        points = np.where(draw_data_obs == 1)
                        if len(points[1]) > 0:
                            x_coords = points[0] + 0.5
                            y_coords = points[1] + 0.5
                            ax.scatter(x_coords, y_coords, color='red', marker='o', s=10, zorder=5)
                ax.invert_yaxis()
                plt.tight_layout()
                # 保存和关闭
                if draw_opt==False:
                    plt.savefig(os.path.join(pic_savepath,f'start_day_distance_81.png'),dpi=900,bbox_inches='tight')
                    plt.close()
                else:
                    plt.show()
                        
                
                
